package class_0822Graph;

import java.util.*;

/**
 * 给个任务，用并查集重新写。
 * @author 90774
 */

public class Coding_Kruskal {
    public static class MySets{
        public HashMap<Node, List<Node>> setMap;
        public MySets(List<Node> nodes){
            for(Node cur :nodes){
                List<Node> set = new ArrayList<Node>();
                set.add(cur);
                setMap.put(cur,set);
            }
        }
        public boolean isSameSet(Node from, Node to){return  setMap.get(from)==setMap.get(to);}
        public void union(Node from, Node to){
            List<Node> fromSet = setMap.get(from);
            List<Node> toSet = setMap.get(to);
            for(Node toNode: toSet){
                fromSet.add(toNode);
                setMap.put(toNode,fromSet);
            }
        }

    }
    public static class EdgeComparator implements Comparator<Edge> {
        @Override
        public int compare(Edge o1,Edge o2){
            return o1.weight - o2.weight;
        }
    }

    public Set<Edge> KruskalMST(Graph graph){
        MySets mySet = new MySets((List<Node>) graph.nodes.values());
        PriorityQueue<Edge> pq = new PriorityQueue<>(new EdgeComparator());
        for(Edge edge:graph.edges){
            pq.offer(edge);
        }
        Set<Edge> result = new HashSet<>();
        while(!pq.isEmpty()){
            Edge edge = pq.poll();
            if(!mySet.isSameSet(edge.from,edge.to)){
                result.add(edge);
                mySet.union(edge.from,edge.to);
            }
        }
        return result;
    }
}
